# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import pickle

from absl import app, flags

import experimental.prophet as prophet
from experimental.prophet.craps.universe import universe_btc
from experimental.prophet.ops.linear_model import linear_regression_predict
from experimental.prophet.ops.fill_query import FillQuery
from coin.experimental.jaewon.research.adaptive_edge import lt_model
from coin.experimental.jaewon.research.adaptive_edge import features
from coin.experimental.jaewon.research.adaptive_edge import y_features
from coin.experimental.jaewon.research.adaptive_edge.util import open_file


def eval_matrix(buy_model, sell_model, date):
  vars = []

  # Timers, sampling points
  ts = prophet.global_timestamp()
  timer_1s = prophet.timer('1s')
  trigger = (ts % 10**9 == 0)
  vars += [ts]

  # Symbols
  universe = universe_btc(date)
  symbol_info = universe[0]
  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs

  # Basic book info, mid-price Y
  vars += y_features.basic(symbol_info, feed)

  # Features
  with prophet.namespace('feature'):
    feature_vars = features.features(date)

  # Sampler
  with prophet.namespace('sampler'):
    vars += features.samplers('feature/')

  # Long-term model theo price
  with prophet.namespace('lt_model'):
    theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, '2m', date=date)
    vars += [theo_ask0_p, theo_bid0_p]

  # Model
  with prophet.namespace('edge_model'):
    x_vars = [v for v in feature_vars if not v.name.endswith('.x')]
    buy_edge_bp = (linear_regression_predict(x_vars, buy_model) + 3.).named_ref('buy_edge')
    sell_edge_bp = (linear_regression_predict(x_vars, sell_model) + 3.).named_ref('sell_edge')
    buy_p = prophet.floor(theo_bid0_p * (1. - buy_edge_bp / 10000.)).named_ref('buy_p')
    sell_p = prophet.ceil(theo_ask0_p * (1. + sell_edge_bp / 10000.)).named_ref('sell_p')
    vars += [buy_edge_bp, sell_edge_bp, buy_p, sell_p]

  # Fill info with model
  with prophet.namespace('fillinfo'):
    fill_query = FillQuery(book, trade, '1s', '8s')
    buy_fill, buy_taker, buy_fill_price = fill_query.query_buy(buy_p)
    sell_fill, sell_taker, sell_fill_price = fill_query.query_sell(sell_p)

    vars += [
        buy_fill.named_ref('buy_filled'),
        buy_taker.named_ref('buy_taker'),
        buy_fill_price.named_ref('buy_fill_price'),
        sell_fill.named_ref('sell_filled'),
        sell_taker.named_ref('sell_taker'),
        sell_fill_price.named_ref('sell_fill_price')
    ]

  with prophet.control_if(trigger):
    aggregator = prophet.scalar_value_aggregator(vars)

  return aggregator


def make_y_inplace(df, y_windows=None):
  y_windows = y_windows or ['1s', '2s', '5s', '10s', '15s', '20s', '30s', '1m', '2m', '5m']

  for w in y_windows:
    y_mid = df['Y/mid_p.Y%s' % w]
    df['Y/buy_ret_%s'
       % w] = (y_mid - df['fillinfo/buy_fill_price']) / y_mid * df['fillinfo/buy_filled']
    df['Y/sell_ret_%s'
       % w] = (df['fillinfo/sell_fill_price'] - y_mid) / y_mid * df['fillinfo/sell_filled']


def main(argv):
  FLAGS = flags.FLAGS

  with open_file(FLAGS.model_in, 'rb') as f:
    model = pickle.load(f)

  buy_model = dict(model['buy']['model'].params)
  sell_model = dict(model['sell']['model'].params)

  df = prophet.run_from_fastfeed_multiprocess(functools.partial(eval_matrix, buy_model, sell_model),
                                              datetime.date(2019, 8, 1),
                                              datetime.date(2019, 8, 10),
                                              machine='feed-01.ap-northeast-1.aws',
                                              max_workers=16,
                                              debug_executor=FLAGS.debug_executor,
                                              inject_date_to_graph_func=True,
                                              duration_after='20m')

  make_y_inplace(df)

  print(df.describe())
  df.to_hdf('dump_eval.hd5.lz4', '_')


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_boolean('debug_executor', False, 'Debug executor')

  flags.DEFINE_string('model_in', 'model.pkl.gz', 'Model output filepath')

  app.run(main)
